package main

import (
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"

	"github.com/dave/jennifer/jen"
	"github.com/stackrox/rox/pkg/set"
	"github.com/stackrox/rox/pkg/utils"
)

var (
	printerRegex = regexp.MustCompile(`(?s)func ([a-zA-z]+)Printer\(`)
)

// Lifted straight from the goimports code
func isGoFile(f os.DirEntry) bool {
	name := f.Name()
	return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
}

func getGoFiles() (fileNames []string, err error) {
	files, err := os.ReadDir(".")
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		if !isGoFile(file) {
			continue
		}
		fileNames = append(fileNames, file.Name())
	}
	return fileNames, nil
}

func processGoFile(fileName string, knownKeys set.StringSet) error {
	f, err := os.Open(fileName)
	if err != nil {
		return err
	}
	defer utils.IgnoreError(f.Close)
	contents, err := io.ReadAll(f)
	if err != nil {
		return err
	}
	for _, match := range printerRegex.FindAllStringSubmatch(string(contents), -1) {
		knownKeys.Add(match[1])
	}
	return nil
}

func capitalizeFirstLetter(k string) string {
	if len(k) == 0 {
		return ""
	}
	return strings.ToUpper(k[0:1]) + k[1:]
}

func main() {
	goFiles, err := getGoFiles()
	utils.CrashOnError(err)
	knownKeys := set.NewStringSet()
	for _, goFile := range goFiles {
		utils.Must(processGoFile(goFile, knownKeys))
	}
	constDefs := make([]jen.Code, 0, len(knownKeys))
	regCalls := make([]jen.Code, 0, len(knownKeys))
	for _, k := range knownKeys.AsSortedSlice(func(i, j string) bool {
		return i < j
	}) {
		constName := fmt.Sprintf("%sKey", capitalizeFirstLetter(k))
		constDefs = append(constDefs, jen.Id(constName).Op("=").Lit(k))
		regCalls = append(regCalls, jen.Id("registerFunc").Call(jen.Id(constName), jen.Id(fmt.Sprintf("%sPrinter", k))))
	}

	f := jen.NewFile(os.Getenv("GOPACKAGE"))
	f.HeaderComment("Code generated by violationprinter codegen. DO NOT EDIT.")

	f.Add(jen.Const().Defs(constDefs...)).
		Line().
		Func().Id("init").Params().Block(regCalls...)
	utils.Must(f.Save("gen-registrations.go"))
}
